from kuai.sim.actions import SimulationAction
import sys 

class SimulationReporter(SimulationAction):
    def __init__(self, output = sys.stdout):
        SimulationAction.__init__(self)
        self.output = output

class SimulationDigest(SimulationReporter):
    def __init__(self, output = sys.stdout):
        SimulationDigest.__init__(self, output)
        self.title = ""
        
    def start(self, mol, efunc):
        self.output.write(self.title)
        self.output.flush()

class DigestPBC(SimulationDigest):
    def __init__(self, output = sys.stdout):
        SimulationDigest.__init__(self, output)
        self.title = """         :        A         B         C      Alpha      Beta     Gamma
         :       (A)       (A)       (A)     (deg)     (deg)     (deg)
---------+------------------------------------------------------------
"""

    def start(self, mol, efunc):
        if mol.pbc:
            SimulationDigest.start(self, mol, efunc)
        else:
            raise RuntimeError("There is NO PBC in this molecule!")
        
    def act(self, mol, istep, efunc):
        format = "%9d:%10.2f%10.2f%10.2f%10.2f%10.2f%10.2f\n"
        message = format % (istep, mol.pbc.a(), mol.pbc.b(), mol.pbc.c(), \
                            mol.pbc.beta(), mol.pbc.gamma(), )
        self.output.write(message)

class DigestMoleculeModel(SimulationDigest):
    def __init__(self, output = sys.stdout):
        SimulationDigest.__init__(self, output)
        self.title = """         :     E Total    Potenial     Kinetic Temperature    Pressure Density
         :    (kJ/mol)    (kJ/mol)    (kJ/mol)         (K)       (MPa)  (g/cc)
---------+--------------------------------------------------------------------
"""

    def act(self, mol, istep, efunc):
        format1 = "%9d:%10.2f%10.2f%10.2f%10.2f%10.2f%10.2f\n"
        message = format % (istep, mol.pbc.a(), mol.pbc.b(), mol.pbc.c(), \
                            mol.pbc.beta(), mol.pbc.gamma(), )
        self.output.write(message)
        
